import pandas as pd
import numpy as np
from datetime import datetime
import akshare as ak
import matplotlib.pyplot as plt
import seaborn as sns
import sys
import os

# 设置中文显示
plt.rcParams["font.sans-serif"] = ['LXGW WenKai Mono Screen',
                                   "SimHei"]  # 用来正常显示中文标签
plt.rcParams["axes.unicode_minus"] = True  # 用来正常显示负号

# 尝试导入颜色常量
try:
    from text import OKGREEN, OKRED, END, BOLD
except ImportError:
    # 如果导入失败，尝试从父目录导入
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    try:
        from text import OKGREEN, OKRED, END, BOLD
    except ImportError:
        # 如果仍然失败，定义默认值
        OKGREEN = '\033[92m'
        OKRED = '\033[91m'
        END = '\033[0m'
        BOLD = '\033[1m'

def get_sh_index_data(start_date: str, end_date: str) -> pd.DataFrame:
    """
    获取上证指数历史数据
    
    Args:
        start_date: 开始日期，格式：YYYY-MM-DD
        end_date: 结束日期，格式：YYYY-MM-DD
        
    Returns:
        DataFrame包含以下列：
        - date: 日期
        - close: 收盘价
        - change_pct: 涨跌幅
    """
    try:
        print("正在获取上证指数数据...")
        
        # 使用akshare获取上证指数数据（不指定日期范围，获取所有可用数据）
        df = ak.stock_zh_index_daily(symbol="sh000001")
        
        if df.empty:
            print("获取的原始数据为空，请检查网络连接。")
            return pd.DataFrame()
            
        print(f"成功获取原始数据，共 {len(df)} 条记录。")
        
        
        # 筛选日期范围
        try:
            # 将开始和结束日期转换为datetime对象
            start_dt = pd.to_datetime(start_date)
            end_dt = pd.to_datetime(end_date)
            
            # 创建一个列表来存储符合条件的行索引
            filtered_indices = []
            
            # 使用for循环手动筛选日期范围内的数据
            for idx, row in df.iterrows():
                row_date = pd.to_datetime(row['date']).date()  # 只比较日期部分，忽略时间
                print(f"正在检查日期 {row_date} 是否在指定日期范围内...")
                if start_dt.date() <= row_date <= end_dt.date():
                    filtered_indices.append(idx)
            
            # 使用索引筛选数据
            df = df.loc[filtered_indices]
            
            if df.empty:
                print(f"在指定日期范围 {start_date} 至 {end_date} 内没有找到数据。")
                return pd.DataFrame()
                
            print(f"日期范围筛选后，剩余 {len(df)} 条记录。")
        except Exception as e:
            print(f"日期筛选出错: {str(e)}")
            return pd.DataFrame()
        
        # 计算涨跌幅
        df['change_pct'] = df['close'].pct_change() * 100
        
        # 删除第一行（没有涨跌幅数据）
        if len(df) > 1:
            df = df.iloc[1:].copy()
            print(f"计算涨跌幅后，最终有效数据 {len(df)} 条记录。")
        else:
            print("数据记录不足，无法计算涨跌幅。")
            return pd.DataFrame()
        
        return df[['date', 'close', 'change_pct']]
    except Exception as e:
        print(f"获取数据失败：{str(e)}")
        print("请检查网络连接是否正常，以及akshare包是否为最新版本。")
        print("您可以尝试运行 'pip install --upgrade akshare' 更新akshare包。")
        return pd.DataFrame()

def analyze_monthly_periods(df: pd.DataFrame) -> dict:
    """
    分析每个月的上旬、中旬、下旬的涨跌幅
    
    Args:
        df: 包含日期和涨跌幅的DataFrame
        
    Returns:
        字典，包含每个月的上中下旬统计数据
    """
    # 添加月份和日期段标签
    df['date'] = pd.to_datetime(df['date'])
    df['month'] = df['date'].dt.strftime('%Y-%m')
    df['day'] = df['date'].dt.day
    
    # 定义日期段
    def get_period(day):
        if day <= 10:
            return '上旬'
        elif day <= 20:
            return '中旬'
        else:
            return '下旬'
    
    df['period'] = df['day'].apply(get_period)
    
    # 按月份和日期段分组计算累计涨跌幅
    results = {}
    for month in df['month'].unique():
        month_data = df[df['month'] == month]
        month_stats = {}
        
        for period in ['上旬', '中旬', '下旬']:
            period_data = month_data[month_data['period'] == period]
            if not period_data.empty:
                # 计算区间累计涨跌幅
                cumulative_return = ((1 + period_data['change_pct']/100).prod() - 1) * 100
                month_stats[period] = {
                    'cumulative_return': cumulative_return,
                    'trading_days': len(period_data),
                    'positive_days': len(period_data[period_data['change_pct'] > 0]),
                    'negative_days': len(period_data[period_data['change_pct'] < 0])
                }
        
        results[month] = month_stats
    
    return results

def print_analysis_results(results: dict):
    """
    打印分析结果
    
    Args:
        results: analyze_monthly_periods函数返回的结果字典
    """
    print(f"\n{BOLD}上证指数月内上中下旬涨跌幅统计{END}\n")
    
    for month, stats in results.items():
        print(f"\n{BOLD}{month}{END}")
        print("-" * 50)
        
        for period in ['上旬', '中旬', '下旬']:
            if period in stats:
                data = stats[period]
                return_color = OKGREEN if data['cumulative_return'] > 0 else OKRED
                print(f"{period}:")
                print(f"  累计涨跌幅: {return_color}{data['cumulative_return']:.2f}%{END}")
                print(f"  交易天数: {data['trading_days']}")
                print(f"  上涨天数: {OKGREEN}{data['positive_days']}{END}")
                print(f"  下跌天数: {OKRED}{data['negative_days']}{END}")
        print("-" * 50)

def results_to_dataframe(results: dict) -> pd.DataFrame:
    """
    将分析结果转换为DataFrame以便可视化
    
    Args:
        results: analyze_monthly_periods函数返回的结果字典
        
    Returns:
        包含分析结果的DataFrame
    """
    data = []
    for month, stats in results.items():
        for period, values in stats.items():
            data.append({
                'month': month,
                'period': period,
                'cumulative_return': values['cumulative_return'],
                'trading_days': values['trading_days'],
                'positive_days': values['positive_days'],
                'negative_days': values['negative_days']
            })
    
    return pd.DataFrame(data)

def plot_monthly_returns(result_df: pd.DataFrame):
    """
    绘制月内上中下旬累计涨跌幅柱状图
    
    Args:
        result_df: 包含分析结果的DataFrame
    """
    plt.figure(figsize=(14, 8))
    # 指定period的顺序为上旬、中旬、下旬
    sns.barplot(x='month', y='cumulative_return', hue='period', data=result_df, hue_order=['上旬', '中旬', '下旬'])
    plt.title('上证指数月内上中下旬累计涨跌幅', fontsize=16, y=1.05)
    plt.xlabel('月份', fontsize=12)
    plt.ylabel('累计涨跌幅(%)', fontsize=12)
    plt.xticks(rotation=45)
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()

def plot_trading_days(result_df: pd.DataFrame):
    """
    绘制上涨天数和下跌天数对比图
    
    Args:
        result_df: 包含分析结果的DataFrame
    """
    plt.figure(figsize=(14, 8))
    
    # 创建分组柱状图
    bar_width = 0.35
    index = np.arange(len(result_df))
    
    plt.bar(index, result_df['positive_days'], bar_width, label='上涨天数', color='red', alpha=0.7)
    plt.bar(index + bar_width, result_df['negative_days'], bar_width, label='下跌天数', color='green', alpha=0.7)
    
    # 添加标签和标题
    plt.xlabel('月份-时期', fontsize=12)
    plt.ylabel('天数', fontsize=12)
    plt.title('上证指数月内上中下旬上涨/下跌天数对比', fontsize=16)
    plt.xticks(index + bar_width / 2, [f"{row['month']}-{row['period']}" for _, row in result_df.iterrows()], rotation=45)
    plt.legend()
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()

def analyze_period_statistics(result_df: pd.DataFrame):
    """
    分析不同时期的统计数据
    
    Args:
        result_df: 包含分析结果的DataFrame
        
    Returns:
        包含统计结果的DataFrame
    """
    # 按时期分组统计
    period_stats = result_df.groupby('period').agg({
        'cumulative_return': ['mean', 'std', 'min', 'max'],
        'trading_days': 'mean',
        'positive_days': 'mean',
        'negative_days': 'mean'
    }).round(2)
    
    # 计算上涨概率
    period_stats['positive_ratio'] = (period_stats[('positive_days', 'mean')] / 
                                     (period_stats[('positive_days', 'mean')] + period_stats[('negative_days', 'mean')])).round(4) * 100
    
    return period_stats

def plot_period_statistics(period_stats):
    """
    绘制各时期统计数据的可视化图表
    
    Args:
        period_stats: 包含时期统计数据的DataFrame
    """
    # 确保索引按照上旬、中旬、下旬的顺序
    ordered_index = ['上旬', '中旬', '下旬']
    period_stats = period_stats.reindex(ordered_index)
    
    # 绘制各时期平均涨跌幅对比图
    plt.figure(figsize=(10, 6))
    sns.barplot(x=period_stats.index, y=period_stats[('cumulative_return', 'mean')], order=ordered_index)
    plt.title('上证指数月内各时期平均涨跌幅对比', fontsize=16, y=1.05)
    plt.xlabel('时期', fontsize=12)
    plt.ylabel('平均涨跌幅(%)', fontsize=12)
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    
    # 在柱状图上添加数值标签
    for i, v in enumerate(period_stats[('cumulative_return', 'mean')]):
        plt.text(i, v + (0.1 if v >= 0 else -0.3), f"{v:.2f}%", ha='center')
    
    plt.tight_layout()
    plt.show()
    
    # 绘制各时期上涨概率对比图
    plt.figure(figsize=(10, 6))
    sns.barplot(x=period_stats.index, y=period_stats['positive_ratio'], order=ordered_index)
    plt.title('上证指数月内各时期上涨概率对比', fontsize=16, y=1.05)
    plt.xlabel('时期', fontsize=12)
    plt.ylabel('上涨概率(%)', fontsize=12)
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    
    # 在柱状图上添加数值标签
    for i, v in enumerate(period_stats['positive_ratio']):
        plt.text(i, v + 1, f"{v:.2f}%", ha='center')
    
    plt.tight_layout()
    plt.show()

def format_date(date_str):
    """
    将日期字符串转换为YYYY-MM-DD格式
    
    Args:
        date_str: 日期字符串，可以是YYYYMMDD或YYYY-MM-DD格式
        
    Returns:
        YYYY-MM-DD格式的日期字符串
    """
    try:
        # 尝试解析YYYYMMDD格式
        date = datetime.strptime(date_str, "%Y%m%d")
    except ValueError:
        try:
            # 尝试解析YYYY-MM-DD格式
            date = datetime.strptime(date_str, "%Y-%m-%d")
        except ValueError:
            print(f"警告: 无效的日期格式 {date_str}，将使用今天的日期")
            date = datetime.now()
    return date.strftime("%Y-%m-%d")

def validate_date(date_str):
    """
    验证日期字符串是否有效且不在未来
    
    Args:
        date_str: 日期字符串
        
    Returns:
        如果日期有效且不在未来，返回YYYY-MM-DD格式的日期字符串；否则返回今天的日期
    """
    try:
        # 首先格式化日期
        formatted_date = format_date(date_str)
        date = datetime.strptime(formatted_date, "%Y-%m-%d")
        
        # 检查日期是否在未来
        now = datetime.now()
        if date > now:
            print(f"{OKRED}警告: 日期 {formatted_date} 在未来！{END}")
            print(f"历史数据分析不能使用未来日期。将使用今天的日期 {OKGREEN}{now.strftime('%Y-%m-%d')}{END} 代替。")
            print("如果您想分析过去的数据，请使用过去的日期范围。")
            return now.strftime("%Y-%m-%d")
        return formatted_date
    except ValueError:
        print(f"{OKRED}警告: 无效的日期格式 {date_str}{END}，将使用今天的日期")
        return datetime.now().strftime("%Y-%m-%d")

def main():
    """
    主函数，用于分析上证指数月内上中下旬涨跌幅
    """
    import argparse
    
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='分析上证指数月内上中下旬涨跌幅')
    parser.add_argument('--start', type=str, default=None, help='开始日期，格式为YYYY-MM-DD或YYYYMMDD，默认为今年年初')
    parser.add_argument('--end', type=str, default=None, help='结束日期，格式为YYYY-MM-DD或YYYYMMDD，默认为今天')
    parser.add_argument('--no-plot', action='store_true', help='不显示图表')
    parser.add_argument('--save-plots', action='store_true', help='保存图表到文件')
    args = parser.parse_args()
    
    # 设置日期范围
    now = datetime.now()
    start_date = args.start if args.start else f"{now.year}-01-01"  # 今年年初
    end_date = args.end if args.end else now.strftime("%Y-%m-%d")  # 今天
    
    # 验证并格式化日期
    start_date = validate_date(start_date)
    end_date = validate_date(end_date)
    
    # 确保开始日期不晚于结束日期
    start_dt = datetime.strptime(start_date, "%Y-%m-%d")
    end_dt = datetime.strptime(end_date, "%Y-%m-%d")
    if start_dt > end_dt:
        print(f"警告: 开始日期 {start_date} 晚于结束日期 {end_date}，将交换这两个日期")
        start_date, end_date = end_date, start_date
    
    print(f"分析时间范围: {start_date} 至 {end_date}")
    
    # 获取数据
    df = get_sh_index_data(start_date, end_date)
    if df.empty:
        print(f"\n{OKRED}未能获取有效数据！{END}")
        print("\n可能的原因:")
        print("1. 指定的日期范围在未来")
        print("2. 指定的日期范围内没有交易数据")
        print("3. 网络连接问题")
        print("\n建议:")
        print(f"- 使用过去的日期，例如: {OKGREEN}{(datetime.now().replace(day=1) - pd.DateOffset(months=6)).strftime('%Y-%m-%d')}{END} 至 {OKGREEN}{datetime.now().strftime('%Y-%m-%d')}{END}")
        print("- 确保日期范围不超过当前日期")
        print("- 检查网络连接是否正常")
        return
    
    # 分析数据
    results = analyze_monthly_periods(df)
    
    # 打印结果
    print_analysis_results(results)
    
    # 将结果转换为DataFrame
    result_df = results_to_dataframe(results)
    
    # 统计分析
    period_stats = analyze_period_statistics(result_df)
    print("\n各时期统计分析:")
    print(period_stats)
    
    # 可视化
    if not args.no_plot:
        print("\n正在生成图表...")
        plot_monthly_returns(result_df)
        plot_trading_days(result_df)
        plot_period_statistics(period_stats)
    
    # 保存图表
    if args.save_plots:
        print("\n正在保存图表...")
        # 创建保存图表的目录
        save_dir = "sh_index_analysis_plots"
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        
        # 保存月内上中下旬累计涨跌幅柱状图
        plt.figure(figsize=(14, 8))
        sns.barplot(x='month', y='cumulative_return', hue='period', data=result_df, hue_order=['上旬', '中旬', '下旬'])
        plt.title('上证指数月内上中下旬累计涨跌幅', fontsize=16, y=1.05)
        plt.xlabel('月份', fontsize=12)
        plt.ylabel('累计涨跌幅(%)', fontsize=12)
        plt.xticks(rotation=45)
        plt.grid(axis='y', linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.savefig(f"{save_dir}/monthly_returns.png")
        plt.close()
        
        # 保存上涨天数和下跌天数对比图
        plt.figure(figsize=(14, 8))
        bar_width = 0.35
        index = np.arange(len(result_df))
        plt.bar(index, result_df['positive_days'], bar_width, label='上涨天数', color='red', alpha=0.7)
        plt.bar(index + bar_width, result_df['negative_days'], bar_width, label='下跌天数', color='green', alpha=0.7)
        plt.xlabel('月份-时期', fontsize=12)
        plt.ylabel('天数', fontsize=12)
        plt.title('上证指数月内上中下旬上涨/下跌天数对比', fontsize=16)
        plt.xticks(index + bar_width / 2, [f"{row['month']}-{row['period']}" for _, row in result_df.iterrows()], rotation=45)
        plt.legend()
        plt.grid(axis='y', linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.savefig(f"{save_dir}/trading_days.png")
        plt.close()
        
        # 确保索引按照上旬、中旬、下旬的顺序
        ordered_index = ['上旬', '中旬', '下旬']
        ordered_period_stats = period_stats.reindex(ordered_index)
        
        # 保存各时期平均涨跌幅对比图
        plt.figure(figsize=(10, 6))
        sns.barplot(x=ordered_period_stats.index, y=ordered_period_stats[('cumulative_return', 'mean')], order=ordered_index)
        plt.title('上证指数月内各时期平均涨跌幅对比', fontsize=16, y=1.05)
        plt.xlabel('时期', fontsize=12)
        plt.ylabel('平均涨跌幅(%)', fontsize=12)
        plt.grid(axis='y', linestyle='--', alpha=0.7)
        for i, v in enumerate(ordered_period_stats[('cumulative_return', 'mean')]):
            plt.text(i, v + (0.1 if v >= 0 else -0.3), f"{v:.2f}%", ha='center')
        plt.tight_layout()
        plt.savefig(f"{save_dir}/period_returns.png")
        plt.close()
        
        # 保存各时期上涨概率对比图
        plt.figure(figsize=(10, 6))
        sns.barplot(x=ordered_period_stats.index, y=ordered_period_stats['positive_ratio'], order=ordered_index)
        plt.title('上证指数月内各时期上涨概率对比', fontsize=16, y=1.05)
        plt.xlabel('时期', fontsize=12)
        plt.ylabel('上涨概率(%)', fontsize=12)
        plt.grid(axis='y', linestyle='--', alpha=0.7)
        for i, v in enumerate(ordered_period_stats['positive_ratio']):
            plt.text(i, v + 1, f"{v:.2f}%", ha='center')
        plt.tight_layout()
        plt.savefig(f"{save_dir}/positive_ratio.png")
        plt.close()
        
        print(f"图表已保存到 {save_dir} 目录")

if __name__ == "__main__":
    main()